
from netCDF4 import Dataset
from numpy.core.fromnumeric import shape, size, transpose
from wrf import getvar
import numpy as np
import math
import os

def list_ncfiles(Dir, ncfiles):
	for f in os.listdir(Dir):
		if f.startswith('wrfout'):
			ncfiles.append(f)
	ncfiles.sort()
	return (ncfiles)
def list_files_8(Dir):
	Post_Processed_Data = []
	for f in os.listdir(Dir):
		Post_Processed_Data.append(f)
	return (Post_Processed_Data)
def create_file (Output_Dir, File_name):
    os.chdir(Output_Dir)
    files = []
    files = list_files_8 (Output_Dir)
    test = False
    for file in files:
        if file == File_name:
            test = True
            break
    if test == False:
        os.mkdir (File_name)
    os.chdir (Output_Dir + File_name)

## Enter full path to directory name Maria
Input_Dir = '/~/Maria/'
## Enter full path to where the output should be stored
Output_Dir = '/~/outputs/'
##post process both PBLs
PBLS = ['YSU','MYJ']
Time_Idx =0
for PBL in PBLS:
        #since YSU has additional file, make the CLS dynamic change
        if PBL=='YSU':
                CLS = ['1.0','km_0_2','km_5_0.0','lvl_4','250']  
        else:
            CLS = ['1.0','km_0_2','km_5_0','lvl_4']   
        for CL in CLS:      
                #specify the name of the dir with files to post process
                Hurricane_Setting = 'Maria_8km_NoTurb_'+ PBL +'_hpbl_'+CL
                Input_Dir_1 = Input_Dir +'/'+PBL+'/'+ Hurricane_Setting
                print(Input_Dir_1)
                print(Hurricane_Setting)
                os.chdir(Input_Dir_1)
                ########## list to hold the wrfout files ##########
                ncfiles = []
                ncfiles = list_ncfiles (Input_Dir_1, ncfiles)
                #list for wspd
                All_Max_WND_SPD_10=[]
                #load each wrfout file and process it
                for ncfile in ncfiles:

                        ncfile = Dataset(ncfile)

                        U10_2D = np.array(getvar(ncfile, "U10", Time_Idx))
                        V10_2D = np.array(getvar(ncfile, "V10", Time_Idx))
                        U10_1D = U10_2D.flatten()
                        V10_1D = V10_2D.flatten()
                        WND_SPD_10 = U10_1D
                        # Calculate the wind intensity at each point of the map.
                        for i in range (WND_SPD_10.size - 1):
                                WND_SPD_10[i] = math.sqrt((U10_1D[i]**2)+(V10_1D[i]**2))
                        # append the maximum wind intensity

                        All_Max_WND_SPD_10.append(np.amax(WND_SPD_10))
                
                print('finished with extracting, now just let me write it into: ', Hurricane_Setting)
                #create a dir to store the csv in
                create_file (Output_Dir,'Maria_'+PBL) 
                Path = os.getcwd()
                MyFile=open('%s.csv' %Hurricane_Setting,'w')
                MyFile.write (Hurricane_Setting+str(Time_Idx) + "\n")
                MyFile.write ("All_Max_WND_SPD_10\n")
                #write the results
                for i in range(len(All_Max_WND_SPD_10)):

                        MyFile.write (str(All_Max_WND_SPD_10[i]) +"\n")

                        
                
